package com.lxg.ai.store;

import cn.hutool.core.collection.CollectionUtil;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import org.mapdb.DB;
import org.mapdb.DBMaker;
import org.mapdb.Serializer;
import org.springframework.util.CollectionUtils;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

import static dev.langchain4j.data.message.ChatMessageDeserializer.messagesFromJson;
import static dev.langchain4j.data.message.ChatMessageSerializer.messagesToJson;

public class PersistentChatMemoryStore implements ChatMemoryStore {

    private final DB db = DBMaker.fileDB("./chat-memory.db").transactionEnable().make();
    private final Map<Integer, String> map = db.hashMap("messages",Serializer.INTEGER, Serializer.STRING).createOrOpen();

    @Override
    public List<ChatMessage> getMessages(Object memoryId) {
        String json = map.get(Integer.valueOf(memoryId.toString()));
        return messagesFromJson(json);
    }

    @Override
    public void updateMessages(Object memoryId, List<ChatMessage> messages) {
        String json = messagesToJson(messages);
        map.put(Integer.valueOf(memoryId.toString()), json);
        db.commit();
    }

    @Override
    public void deleteMessages(Object memoryId) {
        map.remove(Integer.valueOf(memoryId.toString()));
        db.commit();
    }
}
